import uvicorn
import psycopg
from contextlib import asynccontextmanager
from fastapi import FastAPI, HTTPException
import config
from langgraph.store.postgres.aio import AsyncPostgresStore
from langgraph.checkpoint.postgres.aio import AsyncPostgresSaver
import asyncio
from router.store import store_router
from router.chat import chat_router
from router.saver import saver_router





@asynccontextmanager
async def lifespan(app: FastAPI):
    # 初始化 LangGraph 的 AsyncPostgresStore
    conn = await psycopg.AsyncConnection.connect(config.DB_URI, autocommit=True)
    store = AsyncPostgresStore(conn)
    saver = AsyncPostgresSaver(conn)
    await saver.setup()
    await store.setup()

    app.state.store = store
    app.state.saver = saver
    print("✅ AsyncPostgresStore 初始化完成")
    yield
    await conn.close()
    print("🛑 AsyncPostgresStore 关闭")



app = FastAPI(lifespan=lifespan)


app.include_router(store_router, prefix="/api/v1/store", tags=["Store 接口"])
app.include_router(saver_router, prefix="/api/v1/saver", tags=["Saver 接口"])
app.include_router(chat_router, prefix="/api/v1/chat", tags=["Chat 接口"])



if __name__ == "__main__":
    asyncio.set_event_loop_policy(asyncio.WindowsSelectorEventLoopPolicy())
    uvicorn.run("main:app", host="0.0.0.0", port=5050)